package at.chrenko.tu.vs.scheduler;

import java.net.Socket;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Scanner;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.io.IOException;
import java.io.UnsupportedEncodingException;

import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;

import org.bouncycastle.util.encoders.Base64;

import at.chrenko.tu.vs.shared.AESDecorator;
import at.chrenko.tu.vs.shared.Base64Decorator;
import at.chrenko.tu.vs.shared.Channel;
import at.chrenko.tu.vs.shared.Command;
import at.chrenko.tu.vs.shared.EngineState;
import at.chrenko.tu.vs.shared.InputHandler;
import at.chrenko.tu.vs.shared.Message;
import at.chrenko.tu.vs.shared.RSADecorator;
import at.chrenko.tu.vs.shared.SimpleChannel;
import at.chrenko.tu.vs.shared.TaskType;

/**
 * @author Thomas Chrenko - 0728121
 *
 * TCPWorker administers connections to the clients and
 * processes input from them
 *
 */
public class TCPWorker implements Runnable, InputHandler  {

	private Scheduler scheduler;
	private Socket socket;
	private Channel channel;
	private PrivateKey privateKey;
	private PublicKey publicKey;
	private final byte[] schedulerChallenge = new byte[32];
	private final byte[] ivParam = new byte[16];
	final String B64 = "a-zA-Z0-9/+";
	private SecretKey key;
	
	/**
	 * Force client logout and
	 * close the socket => terminate
	 */
	public void closeConnection() {
		try {
			if(socket != null)
				socket.close();
		} catch (IOException e) {}
	}
	
	public TCPWorker(Scheduler scheduler, Socket socket, PrivateKey privateKey, PublicKey publicKey) {
		this.scheduler = scheduler;
		this.socket = socket;
		this.privateKey = privateKey;
		this.publicKey = publicKey;
	}
	
	@Override
	public void processInput(String input) {
		if(input.startsWith(Command.LOGIN)) {
			try {
				String output = handleLogin(input);
				channel.send(output);
				
				channel = new Base64Decorator(new AESDecorator(new Base64Decorator(new SimpleChannel(socket.getInputStream(), socket.getOutputStream())), key, ivParam));
			
				byte[] schedulerChallenge = channel.receive();
				
				if(!Arrays.equals(schedulerChallenge, this.schedulerChallenge)) {
					System.err.println("Error");
					closeConnection();
				}
				
				channel = new AESDecorator(new Base64Decorator(new SimpleChannel(socket.getInputStream(), socket.getOutputStream())), key, ivParam);
				
			} catch (IOException e) {
				System.err.println(e.getMessage());
				closeConnection();
			} catch (NoSuchAlgorithmException e) {
				System.err.println(e.getMessage());
				closeConnection();
			} catch (NoSuchPaddingException e) {
				System.err.println(e.getMessage());
				closeConnection();
			}
		}
		else if(input.startsWith(Command.REQUESTENGINE)) {
			try {
				channel.send(handleRequestEngine(input));
			} catch (IOException e) {
				System.err.println(e.getMessage());
			}//out.println(handleRequestEngine(input));
		}
		else {
			System.out.println(Message.UNKNOWN_COMMAND);
		}
	}
	
	@Override
	public void run() {
		try {
//			in = new BufferedReader(new InputStreamReader(socket.getInputStream()));
//			out = new PrintWriter(socket.getOutputStream(), true);
			channel = new RSADecorator(new Base64Decorator(new SimpleChannel(socket.getInputStream(), socket.getOutputStream())), privateKey, publicKey);
		} catch (IOException e) {
			System.out.println("Error opening input stream!");
			return;
		} catch (NoSuchAlgorithmException e) {
			System.out.println("Error opening input stream!");
			return;
		} catch (NoSuchPaddingException e) {
			System.out.println("Error opening input stream!");
			return;
		}
		
		String input;
		
		try {
			while ((input = channel.receiveAsString()) != null) {//in.readLine()) != null) {
			    processInput(input);
			}
		} catch (IOException e) {}
	}
	
	private String handleLogin(String input) {
		input = input.replaceFirst(Command.LOGIN + " ", "");
		
		StringBuilder response = new StringBuilder(Command.OK);
		response.append(" ");
		
		Scanner scanner = new Scanner(input);
		
		if(!scanner.hasNext()) {
			response.append("ERROR");
			return response.toString();
		}
		
		String managerChallengeBase64 = scanner.next();
		
		if(scanner.hasNext()) {
			response.append("ERROR");
			return response.toString();
		}
		
		SecureRandom secureRandom = new SecureRandom();
		
		secureRandom.nextBytes(schedulerChallenge);
		
		String schedulerChallengeBase64 = "";
		
		try {
			schedulerChallengeBase64 = new String(Base64.encode(schedulerChallenge), "UTF-8");
		} catch (UnsupportedEncodingException e) {
			response.append("ERROR");
			return response.toString();
		}
		
		KeyGenerator generator;
		try {
			generator = KeyGenerator.getInstance("AES");
		} catch (NoSuchAlgorithmException e) {
			response.append("ERROR");
			return response.toString();
		}
		
		generator.init(256);
		
		key = generator.generateKey();
		
		String secretKeyBase64 = "";
		
		try {
			secretKeyBase64 = new String(Base64.encode(key.getEncoded()), "UTF-8");
		} catch (UnsupportedEncodingException e) {
			response.append("ERROR");
			return response.toString();
		}
		
		secureRandom = new SecureRandom();
		
		secureRandom.nextBytes(ivParam);
		
		String ivParamBase64 = "";
		
		try {
			ivParamBase64 = new String(Base64.encode(ivParam), "UTF-8");
		} catch (UnsupportedEncodingException e) {
			response.append("ERROR");
			return response.toString();
		}
		
		try {
			this.channel = new RSADecorator(new Base64Decorator(new SimpleChannel(socket.getInputStream(), socket.getOutputStream())), privateKey, publicKey);
		} catch (NoSuchAlgorithmException e) {
			response.append("ERROR");
			return response.toString();
		} catch (NoSuchPaddingException e) {
			response.append("ERROR");
			return response.toString();
		} catch (IOException e) {
			response.append("ERROR");
			return response.toString();
		}
		
		response.append(managerChallengeBase64);
		response.append(" ");
		response.append(schedulerChallengeBase64);
		response.append(" ");
		response.append(secretKeyBase64);
		response.append(" ");
		response.append(ivParamBase64);
		
		String secondMessage = response.toString();
		assert secondMessage.matches("!ok ["+B64+"]{43}= ["+B64+"]{43}= ["+B64+"]{43}= ["+B64+"]{22}==") : "2nd message";
		
		return secondMessage;
	}

	/**
	 * Handles requestengine command
	 * 
	 * @param input
	 * @return If input is a vlid requestengine command and the client is logged in
	 * 		   it request the load from all engines and returns the engines address and
	 * 		   and an error message otherwise
	 */
	private String handleRequestEngine(String input) {
		input = input.replaceFirst(Command.REQUESTENGINE, "");
		
		StringBuilder response = new StringBuilder(Command.REQUESTENGINE);
		response.append(" ");
		
		Scanner scanner = new Scanner(input);
		
		if(!scanner.hasNext()) {
			response.append(Message.REQUESTENGINE_SCHEDULER_COMMAND_ERROR);
			return response.toString();
		}
		
		TaskType taskType;
		
		try {
			taskType = TaskType.valueOf(TaskType.class, scanner.next());
		} catch (Exception e) {
			response.append(Message.REQUESTENGINE_SCHEDULER_COMMAND_ERROR);
			return response.toString();
		}
		
		if(!scanner.hasNextInt()) {
			response.append(Message.REQUESTENGINE_SCHEDULER_COMMAND_ERROR);
			return response.toString();
		}
		
		int taskID = scanner.nextInt();
		
		if(!scanner.hasNextInt()) {
			response.append(Message.REQUESTENGINE_SCHEDULER_COMMAND_ERROR);
			return response.toString();
		}
		
		int amount = scanner.nextInt();
		
		if(scanner.hasNext()) {
			response.append(Message.REQUESTENGINE_SCHEDULER_COMMAND_ERROR);
			return response.toString();
		}
		
		// request the load from all available engines with RequestLoadCallables
		ExecutorService threadPool = Executors.newCachedThreadPool();
		
		List<RequestLoadCallable> callables = new ArrayList<RequestLoadCallable>();
		
		for(TaskEngine engine : scheduler.getTaskEngines()) {
			if(engine.getState().equals(EngineState.AVAILABLE))
				callables.add(new RequestLoadCallable(scheduler, engine.getAddress(), engine.getTcpPort()));
		}
		
		//invoke all callables and give them 5 seconds to finish
		try {
			threadPool.invokeAll(callables, 5, TimeUnit.SECONDS);
		} catch (InterruptedException e) {
			response.append(Message.REQUESTENGINE_SCHEDULER_COMMAND_ERROR);
			return response.toString();
		}
		
		for(int i=0; i<amount; i++) {
			TaskEngine bestEngine = getBestEngine(taskType);
			
			if(bestEngine == null) {
				response = new StringBuilder(Command.REQUESTENGINE);
				response.append(" ");
				response.append("Error: No engine available for execution. Please try again later.");
				return response.toString();
			}
			
			//update load
			bestEngine.setLoad(bestEngine.getLoad() + taskType.getLoad());
			
			response.append(bestEngine.getAddress().getHostAddress());
			response.append(" ");
			response.append(bestEngine.getTcpPort());
			response.append(" ");
		}
		
		response.append(taskID);
		
		return response.toString();
	}

	/**
	 * Gets the best engine to handle the task
	 * 
	 * @param taskType - The type of the task to handle
	 * @return The best engine to handle the task
	 */
	private TaskEngine getBestEngine(TaskType taskType) {
		TaskEngine retValue = null;
		
		for(TaskEngine engine : scheduler.getTaskEngines())
			if(engine.getState().equals(EngineState.AVAILABLE))
				if((100 - engine.getLoad()) >= taskType.getLoad() && (retValue == null || engine.getMaxConsumption()-engine.getMinConsumption() < retValue.getMaxConsumption()-retValue.getMinConsumption()))
					retValue = engine;
		
		return retValue;
	}
}
